package geppetto.cat.models;


import geppetto.cat.common.Pair;
import geppetto.cat.corpus.BilingualCorpus;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;



public class SparseTranslationTable extends AbstractSparseTranslationTable{

	// Contains the probabilites T(f|e) index by a uniq identifier. The
	// indetifier is fetch using the REF class from the corpus
	public double[] probs;

	// Contains the probability T(f|null)
	public double[] nullProb; // Index by foregin

	BilingualCorpus corpus;

	SparseTranslationTable() {

	}

	public SparseTranslationTable(BilingualCorpus corpus) {
		this.corpus = corpus;
		probs = new double[corpus.ref.getSize()];
		nullProb = new double[corpus.getForeignSize()];
		initializeCoOcurrences();
		// printTable(System.out);
		// common.MyArrays.printDoubleArray(probs, "Probs");
		// common.MyArrays.printDoubleArray(nullProb, "Null");
	}

	public SparseTranslationTable copy() {
		SparseTranslationTable copy = new SparseTranslationTable(this.corpus);
		copy.probs = new double[corpus.ref.getSize()];
		copy.nullProb = new double[corpus.getForeignSize()];
		for (int i = 0; i < corpus.ref.getSize(); i++) {
			copy.probs[i] = this.probs[i];
		}

		for (int i = 0; i < corpus.getForeignSize(); i++) {
			copy.nullProb[i] = this.nullProb[i];
		}

		return copy;
	}

	
	public void checkIfNormalized() {
		System.out.println("Performing sanity chek for translation");
		for (int si = 0; si < corpus.getSourceSize(); si++) {
			double sum = 0;
			for (int fi = 0; fi < corpus.getForeignSize(); fi++) {
				double prob = getProbability(si, fi);
				if (prob < 0) {
					System.out.println("Translation probs negative prob");
					System.exit(-1);
				}
				sum += prob;
			}

			if (sum == 0.0) {
				// Word not seen
				continue;
			}
			if(Math.abs(1 - sum) > 1.0E-8){
				System.out.println(" translation probs " + ": Sum is not 1");
			}
		}
		double sum = 0;
		for (int fi = 0; fi < corpus.getForeignSize(); fi++) {
			sum += getNullProbability(fi);
		}
		if (Math.abs(1 - sum) > 1.0E-8) {
			System.out.println("Trannslation probs do not sum to one");
			System.exit(-1);
		}

		System.out.println("Passed sanity chek for translation");
		System.out.println();
	}
	
	public void clear() {
		Arrays.fill(probs, 0);
		Arrays.fill(nullProb, 0);

	}

	public void saveTable(String directory) {
		try {
			System.out.println("------------Saving Translation Table");
			String outFilename = directory + "translationTablebin.gz";
			DataOutputStream out = new DataOutputStream(new GZIPOutputStream(
					new FileOutputStream(outFilename)));
			out.writeInt(probs.length);
			for (int i = 0; i < probs.length; i++) {
				out.writeDouble(probs[i]);
			}
			out.writeInt(nullProb.length);
			for (int i = 0; i < nullProb.length; i++) {
				out.writeDouble(nullProb[i]);
			}
			out.close();
		} catch (IOException e) {
			e.printStackTrace();
			System.out.println("Error saving Translation Table");

			System.exit(1);
		}
		System.out.println("------------Saving Translation Table END");
	}

	public static SparseTranslationTable LoadTranslationTable(
			BilingualCorpus corpus, String directory) {
		SparseTranslationTable table = new SparseTranslationTable();
		table.corpus = corpus;
		System.out.println("------------Loading Table");
		String inFilename = directory + "/translationTablebin.gz";
		try {
			DataInputStream data_in = new DataInputStream(new GZIPInputStream(
					new FileInputStream(inFilename)));
			int sizeProbs = data_in.readInt();
			table.probs = new double[sizeProbs];
			for (int i = 0; i < sizeProbs; i++) {
				table.probs[i] = data_in.readDouble();
			}
			int sizeNullProbs = data_in.readInt();
			table.nullProb = new double[sizeNullProbs];
			for (int i = 0; i < sizeNullProbs; i++) {
				table.nullProb[i] = data_in.readDouble();
			}
			data_in.close();
		} catch (FileNotFoundException e) {
			System.out.println("Translation Table does not exits " + directory + "translationTablebin.gz" );
			System.exit(1);
		} catch (IOException e) {
			System.out.println("Error reading TranslationTable");
			System.exit(1);
		}
		System.out.println("------------Loading Table END");
		return table;
	}

	public void initializeCoOcurrences() {
		double[] sums = new double[corpus.getSourceSize()];
		Arrays.fill(sums, 0.0);
		Arrays.fill(probs, 0.0);
		Arrays.fill(nullProb, 0);
		double nullSum = 0;
		int sourceSize = corpus.getSourceSize();
		int foreginSize = corpus.getForeignSize();
		// Sum of source values. Position + 1 counts the nulls
		for (int snum = 0; snum < corpus.getNumberOfTrainingSentences(); snum++) {
			if (snum % 5000 == 0) {
				System.out.print(".");
				System.out.flush();
			}
			int[] f = corpus.getForeignSentence(snum,
					BilingualCorpus.TRAIN_CORPUS);
			int[] s = corpus.getSourceSentence(snum,
					BilingualCorpus.TRAIN_CORPUS);
			for (int fi = 0; fi < f.length; fi++) {
				int foreignPos = f[fi];
				// Inc null count
				nullProb[foreignPos]++;
				nullSum++;

				// Increment all other positions
				for (int si = 0; si < s.length; si++) {
					int sourcePos = s[si];
					probs[corpus.ref.getIndex(sourcePos, foreignPos)]++;
					sums[sourcePos]++;
				}
			}	
		}

		for (int i = 0; i < sourceSize; i++) {
			int[] foreingIndex = corpus.ref.getForeignIndices(i);
			for (int j = 0; j < foreingIndex.length; j++) {
				int pos = foreingIndex[j];
				probs[pos] = probs[pos] / sums[i];
			}

		}

		for (int j = 0; j < foreginSize; j++) {
			nullProb[j] = nullProb[j] / nullSum;
		}
		System.out.println("null sum");
		System.out.println(nullSum);
	}

	public void initializeUniform() {
		System.out.println("Initialize uniform");
		for (int i = 0; i < corpus.getSourceSize(); i++) {
			double value = 1.0 / corpus.ref.getNumberIForeign(i);
			int[] positions = corpus.ref.getForeignIndices(i);

			for (int j = 0; j < positions.length; j++) {
				int pos = positions[j];
				// System.out.println("Nr positions " + corpus.ref.getSize() + "
				// Pos " + pos + " " + probs.length);
				probs[pos] = value;
			}
		}
		Arrays.fill(nullProb, 1.0 / nullProb.length);
		// printTable();
	}

	/*
	 * public void printTable(){ corpus.ref.printReferencePos(corpus);
	 * common.MyArrays.printDoubleArray(probs,"probs " + probs.length);
	 * common.MyArrays.printDoubleArray(nullProb,"null" + nullProb.length); }
	 */

	public final double getProbability(int sourceIndex, int foreignIndex) {
		int ref = corpus.ref.getIndex(sourceIndex, foreignIndex);
		if (ref == -1)
			return 0;
		return probs[ref];
	}

	public final double getNullProbability(int foreignIndex) {
		return nullProb[foreignIndex];
	}

	public final void setProbability(int sourceIndex, int foreignIndex,
			double prob) {
		// Sanity check
		assert (!Double.isNaN(prob)) && prob >= 0 && prob < 1.0001 : prob
				+ " is not a probability";
		probs[corpus.ref.getIndex(sourceIndex, foreignIndex)] = prob;

	}

	public final void setNullProbability(int foreignIndex, double prob) {
		// Sanity check
		assert (!Double.isNaN(prob)) && prob >= 0 && prob < 1.0001 : prob
				+ " is not a probability";
		nullProb[foreignIndex] = prob;
	}

	public void setProbabilityByPos(int pos, double prob) {
		// System.out.println("Changing pos " + pos);
		probs[pos] = prob;

	}

	public boolean equals(Object table) {
		if (!(table instanceof SparseTranslationTable))
			return false;
		for (int i = 0; i < probs.length; i++) {
			if (probs[i] != ((SparseTranslationTable) table).probs[i])
				return false;
		}
		for (int i = 0; i < nullProb.length; i++) {
			if (nullProb[i] != ((SparseTranslationTable) table).nullProb[i])
				return false;
		}
		return true;
	}

	public double[] getProbabilitiesForSource(int sourceId) {
		int[] fPos = corpus.ref.getForeignIndices(sourceId);
		double[] p = new double[fPos.length];
		for (int i = 0; i < p.length; i++) {
			p[i] = probs[fPos[i]];
		}
		return p;
	}

	
	class CompareDic implements Comparator{

		public int compare(Object arg0, Object arg1) {
			if(arg0 instanceof Pair && arg1 instanceof Pair){
				if(((Double)((Pair)arg0).second()) > ((Double)((Pair)arg1).second())){
					return -1;
				}else{
					return 1;
				}
			}else{
				return 0;
			}
		}
		
	}
	
	public ArrayList<Pair> getForeingWords(int sourceId){
		ArrayList<Pair> dic = new ArrayList<Pair>();
		int foreginSize = corpus.getForeignSize();
		for (int foreingIndex = 0; foreingIndex < foreginSize; foreingIndex++) {
			int ref = corpus.ref.getIndex(sourceId, foreingIndex);
			if (ref == -1)
				continue;
			double prob = getProbability(sourceId, foreingIndex);
			if(prob > 0){
				Pair p = new Pair(foreingIndex, prob);
				dic.add(p);
			}
		}
		Collections.sort(dic,new CompareDic());
		//Sort dictionary
		return dic;
	}
	
	
	public void printTable(PrintStream out) {

		int sourceSize = corpus.getSourceSize();
		int foreginSize = corpus.getForeignSize();
		System.out.println("foreign word || source word || prob");
		for (int sourceIndex = 0; sourceIndex < sourceSize; sourceIndex++) {
			for (int foreingIndex = 0; foreingIndex < foreginSize; foreingIndex++) {	
				int ref = corpus.ref.getIndex(sourceIndex, foreingIndex);
				if (ref == -1)
					continue;
				double prob = probs[ref];
				if (prob > 0.0001) {
					out.println(corpus.getForeignWordById(foreingIndex) + "-" + foreingIndex
							+ " | " + corpus.getSourceWordById(sourceIndex) + "-" + sourceIndex
							+ " : " + prob + " ");
				}
			}	
		}
		for (int foreingIndex = 0; foreingIndex < foreginSize; foreingIndex++) {
			if (nullProb[foreingIndex] > 0) {
				out.println( corpus.getForeignWordById(foreingIndex) + "-" + foreingIndex + " || Null  : " + nullProb[foreingIndex] + " ");
			}
		}
	}

	
	public void printSimpleTable(PrintStream out) {

		int sourceSize = corpus.getSourceSize();
		int foreginSize = corpus.getForeignSize();
		for (int sourceIndex = 0; sourceIndex < sourceSize; sourceIndex++) {
			for (int foreingIndex = 0; foreingIndex < foreginSize; foreingIndex++) {	
				int ref = corpus.ref.getIndex(sourceIndex, foreingIndex);
				if (ref == -1)
					continue;
				double prob = probs[ref];
				out.println(corpus.getForeignWordById(foreingIndex) 
							+ "\t" + corpus.getSourceWordById(sourceIndex) 
							+ "\t" + prob + " ");
			}	
		}
		for (int foreingIndex = 0; foreingIndex < foreginSize; foreingIndex++) {
			if (nullProb[foreingIndex] > 0) {
				out.println( corpus.getForeignWordById(foreingIndex) + "\tNull\t" + nullProb[foreingIndex] + " ");
			}
		}
	}
	
	
	public void printTableAll(PrintStream out) {

		int sourceSize = corpus.getSourceSize();
		int foreginSize = corpus.getForeignSize();
		for (int i = 0; i < foreginSize; i++) {
			for (int j = 0; j < sourceSize; j++) {
				int ref = corpus.ref.getIndex(j, i);
				if (ref == -1)
					continue;
				double prob = probs[ref];
				out.println("P(" + corpus.getForeignWordById(i) + " | "
						+ corpus.getSourceWordById(j) + ") " + prob + " ");

			}

			out.println("P(" + corpus.getForeignWordById(i) + " | " + "Null ) "
					+ nullProb[i] + " ");

		}

	}

}
